#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import json
import cv2
import numpy as np
import shutil
import random
from pathlib import Path
import argparse


class PolygonMerger:
    """
    PolygonMerger
    ============

    功能：
        将Label Studio项目中的果实和果梗多边形按关系融合为单一标注，
        并生成YOLO格式的标注文件。

    输入：
        - project_path: 项目路径（包含images/、notes.json）
        - json_export_path: Label Studio导出的JSON文件路径

    输出：
        - labels/: YOLO格式标注文件
        - visualization/: 可视化结果（可选）

    使用方法：
        merger = PolygonMerger(project_path, json_export_path)
        merger()
    """

    def __init__(self, project_path, json_export_path, seed=42):
        self.project_path = Path(project_path)
        self.json_export_path = Path(json_export_path)
        self.seed = seed

        random.seed(seed)
        np.random.seed(seed)

        self.images_dir = self.project_path / "images"
        self.labels_dir = self.project_path / "labels"
        self.notes_path = self.project_path / "notes.json"

        # 创建输出目录
        if self.labels_dir.exists():
            shutil.rmtree(self.labels_dir)
        os.makedirs(self.labels_dir, exist_ok=True)

        # 读取标签映射
        self.name_to_yolo_id, self.id_to_name = self._load_label_mappings()

        # 读取Label Studio的JSON导出文件
        with open(self.json_export_path, 'r', encoding='utf-8') as f:
            self.json_data = json.load(f)

        # 获取所有图片文件
        self.all_image_files = sorted([
            f for f in os.listdir(self.images_dir)
            if f.lower().endswith(('.png', '.jpg', '.jpeg'))
        ])

    def _load_label_mappings(self):
        """从notes.json中读取标签映射"""
        if not self.notes_path.exists():
            raise FileNotFoundError(f"未找到 {self.notes_path}")

        with open(self.notes_path, 'r', encoding='utf-8') as f:
            notes_data = json.load(f)

        name_to_yolo_id = {
            category["name"]: category["id"]
            for category in notes_data["categories"]
        }
        id_to_name = {
            category["id"]: category["name"]
            for category in notes_data["categories"]
        }

        return name_to_yolo_id, id_to_name

    def _extract_filename(self, path):
        """从路径中提取文件名"""
        return os.path.basename(path)

    def _generate_random_color(self):
        """生成随机颜色"""
        return (
            random.randint(50, 255),
            random.randint(50, 255),
            random.randint(50, 255)
        )

    def _process_image_annotations(self, task, image_file):
        """处理图像标注数据"""
        image_annotations = []
        image_relations = []

        task_image_path = task.get('data', {}).get('image', '')
        task_filename = self._extract_filename(task_image_path)

        if task_filename.lower() == image_file.lower():
            annotations = task.get('annotations', [])
            if annotations:
                latest_annotation = annotations[0]
                results = latest_annotation.get('result', [])

                # 分离标注和关系
                for item in results:
                    if item.get('type') == 'relation':
                        image_relations.append(item)
                    elif 'polygonlabels' in item.get('type', ''):
                        image_annotations.append(item)

        return image_annotations, image_relations

    def _convert_points_to_pixels(self, points, img_w, img_h):
        """转换百分比坐标为像素坐标"""
        pixel_points = []
        for x, y in points:
            px = int(x * img_w / 100)
            py = int(y * img_h / 100)
            pixel_points.append((px, py))
        return pixel_points

    def _merge_polygons(self, points1, points2, img_h, img_w, kernel_size, dilate_iterations, erode_iterations):
        """使用膨胀腐蚀方法合并多边形"""
        try:
            # 创建一个空白掩码
            mask = np.zeros((img_h, img_w), dtype=np.uint8)

            # 绘制第一个多边形（果梗）
            poly1_points = np.array(points1, dtype=np.int32)
            cv2.fillPoly(mask, [poly1_points], 255)

            # 绘制第二个多边形（果实）
            poly2_points = np.array(points2, dtype=np.int32)
            cv2.fillPoly(mask, [poly2_points], 255)

            # 应用膨胀操作连接接近的多边形
            kernel = np.ones((kernel_size, kernel_size), np.uint8)
            dilated_mask = cv2.dilate(mask, kernel, iterations=dilate_iterations)

            # 应用腐蚀操作恢复大致形状
            eroded_mask = cv2.erode(dilated_mask, kernel, iterations=erode_iterations)

            # 从掩码中提取轮廓
            contours, _ = cv2.findContours(eroded_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            if contours:
                # 找到最大的轮廓
                largest_contour = max(contours, key=cv2.contourArea)

                epsilon = 0.0001 * cv2.arcLength(largest_contour, True)
                approx = cv2.approxPolyDP(largest_contour, epsilon, True)

                # 转换为点列表
                if len(approx) >= 3:
                    merged_points = [(point[0][0], point[0][1]) for point in approx]
                    return merged_points
        except Exception as morph_error:
            print(f"形态学操作失败: {morph_error}")

        return None

    def _points_to_yolo_format(self, points, img_w, img_h):
        """将像素坐标转换为归一化坐标"""
        normalized_points = []
        for x, y in points:
            nx = x / img_w
            ny = y / img_h
            normalized_points.extend([nx, ny])
        return normalized_points

    def _extract_polygons_from_mask(self, color_mask, img_w, img_h, color_to_class_id, tolerance=3):
        """使用颜色容差提取多边形，解决绘制精度问题"""
        yolo_annotations = []

        # 统计所有找到的颜色
        unique_colors = np.unique(color_mask.reshape(-1, 3), axis=0)
        unique_colors = [tuple(color) for color in unique_colors if not np.array_equal(color, [0, 0, 0])]

        # 为每个目标颜色创建容差范围
        for target_color, class_id in color_to_class_id.items():
            # 创建颜色范围掩码
            lower_bound = np.array([
                max(0, target_color[0] - tolerance),
                max(0, target_color[1] - tolerance),
                max(0, target_color[2] - tolerance)
            ])
            upper_bound = np.array([
                min(255, target_color[0] + tolerance),
                min(255, target_color[1] + tolerance),
                min(255, target_color[2] + tolerance)
            ])

            mask = cv2.inRange(color_mask, lower_bound, upper_bound)

            # 查找轮廓
            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            all_points = []
            for contour in contours:
                # 忽略太小的轮廓
                if cv2.contourArea(contour) < 10:
                    continue

                # 多边形近似
                epsilon = 0.002 * cv2.arcLength(contour, True)
                approx = cv2.approxPolyDP(contour, epsilon, True)

                if len(approx) < 3:
                    continue

                # 归一化坐标
                for point in approx:
                    x = point[0][0] / img_w
                    y = point[0][1] / img_h
                    all_points.append(x)
                    all_points.append(y)

            # 如果有有效点，添加到标注
            if all_points:
                line = f"{class_id} " + " ".join([f"{p:.6f}" for p in all_points])
                yolo_annotations.append(line)

        # 报告未映射的颜色（用于调试）
        unmapped_count = 0
        for color in unique_colors:
            mapped = False
            for target_color in color_to_class_id.keys():
                if (abs(color[0] - target_color[0]) <= tolerance and
                        abs(color[1] - target_color[1]) <= tolerance and
                        abs(color[2] - target_color[2]) <= tolerance):
                    mapped = True
                    break
            if not mapped:
                unmapped_count += 1

        if unmapped_count > 0:
            print(f"警告: 发现 {unmapped_count} 个未映射的颜色区域（已通过容差处理）")

        return yolo_annotations

    def process_images(self):
        """处理所有图像"""
        for idx, image_file in enumerate(self.all_image_files):
            # 读取图像
            img_path = self.images_dir / image_file
            img = cv2.imread(str(img_path))

            if img is None:
                print(f"错误: 无法读取图像 {image_file}")
                continue

            img_h, img_w = img.shape[:2]

            # 根据图像格式设置不同的膨胀腐蚀参数
            if image_file.lower().endswith('.png'):
                # PNG格式 (384x384) - 使用较小的参数
                dilate_iterations = 2
                erode_iterations = 2
                kernel_size = 3
            else:
                # JPG格式 (4624x3472) - 使用较大的参数
                dilate_iterations = 2
                erode_iterations = 2
                kernel_size = 5

            # 创建对应的txt文件路径
            txt_filename = img_path.stem + '.txt'
            txt_path = self.labels_dir / txt_filename

            # 在JSON数据中查找当前图像的标注
            image_annotations = []
            image_relations = []

            # 查找JSON中对应的任务
            found_match = False
            for task in self.json_data:
                image_annotations, image_relations = self._process_image_annotations(task, image_file)
                if image_annotations or image_relations:
                    found_match = True
                    break

            if not found_match or not image_annotations:
                print(f"警告: 未找到图像 {image_file} 的标注")
                continue

            # 创建存储原始坐标点和标签的字典
            id_to_points = {}
            id_to_label = {}

            # 处理所有多边形标注
            for annotation in image_annotations:
                annotation_id = annotation.get('id')
                value = annotation.get('value', {})

                if 'points' in value and 'polygonlabels' in value:
                    # 转换百分比坐标为像素坐标
                    points = self._convert_points_to_pixels(value['points'], img_w, img_h)

                    # 存储原始坐标点和标签
                    if len(points) >= 3:
                        id_to_points[annotation_id] = points
                        id_to_label[annotation_id] = value['polygonlabels'][0]
                    else:
                        print(f"警告: 多边形 {annotation_id} 点数不足 ({len(points)}), 跳过")

            # 如果没有关系，跳过此图像
            if not image_relations:
                print(f"警告: 未找到图像 {image_file} 的关系，跳过")
                continue

            # 存储果梗和果实的对应关系
            stem_to_fruit = {}

            # 收集所有参与关系的多边形ID
            related_polygon_ids = set()

            # 首先处理所有关系
            for relation in image_relations:
                from_id = relation.get('from_id')  # 果梗
                to_id = relation.get('to_id')  # 果实

                # 记录参与关系的多边形ID
                related_polygon_ids.add(from_id)
                related_polygon_ids.add(to_id)

                # 记录果梗和果实的对应关系
                stem_to_fruit[from_id] = to_id

            # 处理关联关系 - 使用膨胀腐蚀方法合并多边形
            merged_polygons = {}
            processed_pairs = set()

            for relation in image_relations:
                from_id = relation.get('from_id')  # 果梗
                to_id = relation.get('to_id')  # 果实

                # 确保我们不会重复处理同一对多边形
                pair_key = tuple(sorted([from_id, to_id]))
                if pair_key in processed_pairs:
                    continue

                processed_pairs.add(pair_key)

                if from_id in id_to_points and to_id in id_to_points:
                    points1 = id_to_points[from_id]  # 果梗
                    points2 = id_to_points[to_id]  # 果实

                    merged_points = self._merge_polygons(
                        points1, points2, img_h, img_w, kernel_size,
                        dilate_iterations, erode_iterations
                    )

                    if merged_points:
                        # 存储合并结果 - 使用果实的标签
                        merged_polygons[from_id] = (merged_points, to_id)
                        merged_polygons[to_id] = (merged_points, to_id)

            # 准备YOLO格式的标注数据
            yolo_annotations = []

            # 只处理参与关系的多边形
            processed_merged = set()
            for annotation_id in related_polygon_ids:
                if annotation_id not in id_to_points:
                    continue

                # 检查这个多边形是否被合并过
                if annotation_id in merged_polygons:
                    if annotation_id in processed_merged:
                        continue

                    merged_points, fruit_id = merged_polygons[annotation_id]

                    # 获取果实的标签
                    if fruit_id in id_to_label:
                        label_name = id_to_label[fruit_id]
                    else:
                        label_name = id_to_label.get(annotation_id, "unknown")

                    if label_name not in self.name_to_yolo_id:
                        print(f"警告: 未知标签 '{label_name}', 跳过")
                        continue

                    class_id = self.name_to_yolo_id[label_name]

                    # 将像素坐标转换为归一化坐标
                    normalized_points = self._points_to_yolo_format(merged_points, img_w, img_h)

                    # 添加到YOLO标注列表
                    yolo_annotations.append(f"{class_id} " + " ".join([f"{p:.6f}" for p in normalized_points]))

                    # 标记这个合并多边形已处理
                    for mid in merged_polygons:
                        if merged_polygons[mid][0] == merged_points:
                            processed_merged.add(mid)
                else:
                    # 如果多边形没有合并，但有关系，则保留原始形状
                    # 只处理果实，果梗将单独处理
                    if annotation_id in stem_to_fruit:  # 这是果梗
                        continue  # 跳过果梗，它们将单独处理

                    points = id_to_points[annotation_id]
                    label_name = id_to_label.get(annotation_id, "unknown")

                    if label_name not in self.name_to_yolo_id:
                        print(f"警告: 未知标签 '{label_name}', 跳过")
                        continue

                    class_id = self.name_to_yolo_id[label_name]

                    # 将像素坐标转换为归一化坐标
                    normalized_points = self._points_to_yolo_format(points, img_w, img_h)

                    # 添加到YOLO标注列表
                    yolo_annotations.append(f"{class_id} " + " ".join([f"{p:.6f}" for p in normalized_points]))

            # 创建彩色掩码用于最终的多边形提取
            color_mask = np.zeros((img_h, img_w, 3), dtype=np.uint8)

            # 创建颜色到类别ID的映射
            color_to_class_id = {}

            # 为每个实例创建唯一ID到颜色的映射
            instance_to_color = {}
            instance_to_class_id = {}

            # 存储多边形数据
            polygons = []
            stem_polygons = []  # 专门存储果梗多边形

            # 处理每个relation，按照相同的顺序
            for relation in image_relations:
                from_id = relation.get('from_id')  # 果梗
                to_id = relation.get('to_id')  # 果实

                # 为每个果实实例生成唯一颜色
                if to_id not in instance_to_color:
                    # 获取类别ID
                    if to_id in id_to_label:
                        label_name = id_to_label[to_id]
                        if label_name in self.name_to_yolo_id:
                            class_id = self.name_to_yolo_id[label_name]
                            class_name = self.id_to_name.get(class_id, "unknown")

                            # 生成唯一颜色
                            color = self._generate_random_color()

                            # 存储实例到颜色的映射
                            instance_to_color[to_id] = color
                            instance_to_class_id[to_id] = class_id

                            # 存储颜色到类别ID的映射
                            color_to_class_id[tuple(color)] = class_id

                # 处理果实多边形
                if to_id in merged_polygons:
                    # 处理合并后的多边形
                    merged_points, _ = merged_polygons[to_id]

                    # 获取颜色和类别
                    color = instance_to_color.get(to_id, (255, 0, 0))
                    class_id = instance_to_class_id.get(to_id, 0)

                    # 添加到多边形列表
                    polygons.append({
                        'class_id': class_id,
                        'points': merged_points,
                        'color': color
                    })
                elif to_id in id_to_points:
                    # 处理未合并的果实多边形
                    points = id_to_points[to_id]

                    # 获取颜色和类别
                    color = instance_to_color.get(to_id, (255, 0, 0))
                    class_id = instance_to_class_id.get(to_id, 0)

                    # 添加到多边形列表
                    polygons.append({
                        'class_id': class_id,
                        'points': points,
                        'color': color
                    })

                # 处理果梗多边形
                if from_id in id_to_points:
                    points = id_to_points[from_id]

                    # 使用对应果实的颜色
                    fruit_id = stem_to_fruit.get(from_id)
                    if fruit_id and fruit_id in instance_to_color:
                        color = instance_to_color[fruit_id]

                        # 添加到果梗多边形列表
                        stem_polygons.append({
                            'points': points,
                            'color': color
                        })

            # 先绘制非果梗多边形（果实）
            for polygon in polygons:
                points = polygon['points']
                color = polygon['color']

                # 在彩色掩码上填充多边形
                pts = np.array(points, np.int32)
                cv2.fillPoly(color_mask, [pts], color)

            # 最后绘制果梗多边形（确保在最上层）
            for polygon in stem_polygons:
                points = polygon['points']
                color = polygon['color']

                pts = np.array(points, np.int32)
                # 在彩色掩码上填充多边形（覆盖之前的内容）
                cv2.fillPoly(color_mask, [pts], color)

            # 从彩色掩码图中提取多边形区域
            yolo_annotations_from_mask = self._extract_polygons_from_mask(
                color_mask, img_w, img_h, color_to_class_id, tolerance=3
            )

            # 使用从掩码提取的标注替换原始标注
            yolo_annotations = yolo_annotations_from_mask

            # 写入YOLO格式的标注文件
            with open(txt_path, 'w', encoding='utf-8') as f:
                for annotation in yolo_annotations:
                    f.write(annotation + '\n')

            print(f"已处理 {image_file}: {len(yolo_annotations)} 个标注写入 {txt_filename}")

    def __call__(self):
        print(f"固定随机种子为 {self.seed}")
        print("开始融合果实与果梗多边形 ...")
        self.process_images()
        print("任务完成")


def main():
    parser = argparse.ArgumentParser(description='果实与果梗融合工具')
    parser.add_argument('project_path', type=str, help='项目路径（包含images/、notes.json）')
    parser.add_argument('json_export_path', type=str, help='Label Studio导出的JSON文件路径')

    args = parser.parse_args()

    merger = PolygonMerger(args.project_path, args.json_export_path)
    merger()


if __name__ == "__main__":
    main()